/*
  Copyright 2010-2011 Patrik Jonsson, sunrise@familjenjonsson.org

  This file is part of Sunrise.

  Sunrise is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  Sunrise is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.
  
*/

#include<iostream>
#include "CCfits/CCfits"
#include "counter.h"
#include "optical.h"
#include "scatterer.h"
#include "constants.h"
#include "dust_model.h"
#include "emission.h"
#include "emission-fits.h"
#include "full_sed_emergence.h"
#include "full_sed_grid.h"
#include "emergence-fits.h"
#include "mcrx-units.h"
#include "grid_factory.h"
#include "shoot.h"
#include "fits-utilities.h"
#include "boost/shared_ptr.hpp"
#include "density_generator.h"
#include "preferences.h"
#include "model.h"
#include "baes_grain_model.h"
#include "xfer.h"
#include "ir_grid.h"
#include "equilibrium.h"
#include "vecops.h"
#include "grain_model_hack.h"
#include "grid-fits.h"

//using namespace blitz;
using namespace mcrx;
using namespace CCfits;
using namespace std; 
using constants::pi;

typedef local_random T_rng_policy;
typedef scatterer<polychromatic_scatterer_policy, T_rng_policy> T_scatterer;
typedef T_scatterer::T_lambda T_lambda;
typedef T_scatterer::T_biaser T_biaser;
typedef dust_model<T_scatterer, cumulative_sampling, T_rng_policy> T_dust_model;
typedef full_sed_emergence T_emergence;

typedef full_sed_grid<adaptive_grid> T_grid;
typedef ir_grid<adaptive_grid, cumulative_sampling, local_random> T_ir_grid;

class stellar_grid : 
  public adaptive_grid<grid_cell_emission<polychromatic_policy, 
					  T_rng_policy> >/*, 
  public emission_collection<polychromatic_policy, cumulative_sampling, 
  T_rng_policy>*/ {
public:
  typedef grid_cell_emission<polychromatic_policy,
			     T_rng_policy > T_emitter;
  typedef adaptive_grid<T_emitter> T_grid;
  typedef T_emitter::T_lambda T_content;
  
  stellar_grid (const vec3d & mi, 
		const vec3d & ma, 
		const vec3i & nn,
		grid_factory<T_emitter>& f) :
    adaptive_grid<T_emitter>(mi, ma, nn, f) {};
};


template<typename T> class datamaker {};

template<> class datamaker<T_grid::T_data> {
public:
  static const array_1& stuff(T_grid::T_data* d) {
    return d->densities();};

  static T_grid::T_data get_data(const T_densities& rho, int nl, 
				 const T_grid::T_cell* c) {
    return T_grid::T_data(rho, vec3d(0,0,0), T_lambda(nl));
  };
};

template<> class datamaker<stellar_grid::T_data> {
public:
  static const array_1& stuff(stellar_grid::T_data* d) {
    return d->get_emission();};

  static stellar_grid::T_data get_data(const T_lambda& em, int nl, 
				       const stellar_grid::T_cell* c) {
    assert(c->is_leaf());
    //note that we need to multiply by the volume to get the total
    //luminosity in the cell, since em is the specific luminosity (per
    //volume).
    return stellar_grid::T_data(T_lambda(em/* *c->volume()*/), *c);
  };
};



template<typename grid_type>
class validation_factory: public grid_factory<typename grid_type::T_data> {
  typedef typename grid_factory<typename grid_type::T_data>::T_grid T_grid;
  typedef typename grid_factory<typename grid_type::T_data>::T_cell T_cell;
  typedef typename grid_factory<typename grid_type::T_data>::T_data T_data;

public:
  const T_densities rho0_, n_tol_, stdev_;
  const T_float hr_, hz_, rmax_, zmax_;
  const int n_lambda, maxlevel;

  T_densities get_rho(const T_cell& c) {
    // return 0 if outside of radius rmax
    const T_float r = sqrt(c.getcenter()[0]*c.getcenter()[0]+
			   c.getcenter()[1]*c.getcenter()[1]);
    const T_float z = abs(c.getcenter()[2]);

    T_densities rho = ((r < rmax_) && (z < zmax_)) ?
      array_1(rho0_*exp(-r/hr_ - z/hz_)) : T_densities(rho0_*0.);
    return rho;
  };

    validation_factory (const T_densities& rho0,
			T_float hr, T_float hz, T_float rmax, T_float zmax,
			int nl, const T_densities& n_tol, 
			const T_densities& stdev, int ml):
      rho0_ (rho0), hr_(hr), hz_(hz), rmax_(rmax), zmax_(zmax),
      n_lambda (nl), maxlevel(ml), stdev_(stdev), n_tol_(n_tol) {};
  
  virtual bool refine_cell_p (const T_cell& c, int level) {
    return (level<maxlevel);
      
    // cell is refined to keep column density below n_tol
    T_densities rho = get_rho(c);
    T_densities n2 (dot(c.getsize(), c.getsize())*rho*rho);
    bool refine = all( n2 > (n_tol_*n_tol_));
    //if(!refine) cout << "No refinement level " << level << " at " << c.getcenter() << c.getsize() << endl;
    return refine;
  };

  virtual bool unify_cell_p (typename T_grid::const_local_iterator begin,
			     typename T_grid::const_local_iterator end) {
    T_densities s (rho0_*0.);
    T_densities ssq (rho0_*0.);
    int n=1;
    while (begin != end) {
      T_densities o = datamaker<T_data>::stuff((begin++)->data());
      s+=o;
      ssq+=o*o;
      n++;
    }

    const T_densities mean (s/n);
    const T_densities stdev2 (ssq/n- mean*mean);
    return all ((mean==0) || ((stdev2/(mean*mean)<stdev_*stdev_) && (mean<n_tol_)));
  };

  virtual T_data get_data (const T_cell& c) {
    return datamaker<T_data>::get_data(get_rho(c), n_lambda, &c);
  }
    
  virtual int n_threads () {return 1;};
  virtual int work_chunk_levels () {return 5;};
  virtual int estimated_levels_needed () {return 0;};
};


void baes_validation (T_float tau_f, const string& output_file)
{
  const int n_threads = 12;
  const long n_rays =10000000;
  const int n_rays_temp = 10000000;
  const int maxlevel=4;
  const int n_forced=1;
  const int ir_ref_lambda = 33;
  const T_float rel_tol =1e-2;
  const T_float abs_tol =1e-2;
  const T_float tau_tol = 0.25;
  T_densities stdev_tol(1); stdev_tol = 0.0;
  const vec3d grid_extent(12,12,1);
  const vec3i grid_n(10,10,10);
  const T_float hrd=3.0;
  const T_float hzd=0.15;
  const T_float rdmax=12.0;
  const T_float zdmax=0.9;
  //const T_float tau_f=1.0;

  const T_float hrs=3.0;
  const T_float hzs=0.30;
  const T_float rsmax=12.0;
  const T_float zsmax=1.8;
  const T_float L_tot=1.919e37;

  const int npix=501;
  const T_float cam_fov=24;
  const T_float cam_dist=10e3;

  mcrx::seed(42);
  T_unit_map units;
  units ["length"] = "kpc";
  units ["mass"] = "Msun";
  units ["luminosity"] = "W";
  units ["wavelength"] = "m";
  units ["L_lambda"] = "W/m";

  // create wavelengths
  array_1 lambda(logspace(1e-7,1e-3,81));
  vector<T_float> lambdas(lambda.begin(), lambda.end());

  Preferences p;
  //p.setValue("grain_data_directory", string("/home/patrik/dust_data/crosssections"));
  p.setValue("grain_data_directory", string("/n/circelfs/hernquist_lab/pjonsson/dust_data/crosssections"));
  p.setValue("n_threads", n_threads);
  //p.setValue("use_grain_temp_lookup", true);
  //p.setValue("use_cuda", true);

  // load grain model
  T_dust_model::T_scatterer_vector sv;
  sv.push_back(boost::shared_ptr<T_scatterer> 
	       //(new wd01_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>(p, units)));
	       (new baes_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>(p.getValue("grain_data_directory", string()), p, units)));

  T_dust_model model (sv);
  model.set_wavelength(lambda);

  // we set this manually
  const int lambda1um=20;
  //const int lambda1um = lower_bound(lambdas.begin(), lambdas.end(),
  //			    1e-6) - lambdas.begin()-1;

  //cout << "Lambda 1um is " << lambda1um << '\t' << lambdas[lambda1um] << endl;
  const int vis_ref_lambda = lower_bound(lambdas.begin(), lambdas.end(),
					 .9e-6) - lambdas.begin()-1;

  // find opacity normalization for whatever dust we're using
  const T_float unit_opacity_rho = 1.0/sv[0]->opacity()(lambda1um);
  T_densities rho0(1); 
  // calculate central density from face-on optical depth
  rho0 = tau_f*unit_opacity_rho/(2*hzd*(1-exp(-zdmax/hzd)));
  cout << "Unit optical depth column density is " << unit_opacity_rho << endl;
  cout << "Face-on optical depth is " << tau_f << endl;
  cout << "Central density is " << rho0 << endl;
  T_densities n_tol(1); n_tol = tau_tol*unit_opacity_rho;
  stdev_tol = stdev_tol;
  cout << "Stdev tolerance is " << stdev_tol << endl;
  cout << "Column density tolerance is " << n_tol << endl;

  // build dust grid
  cout << "\nBuilding dust grid" << endl;
  validation_factory<T_grid> factory (rho0, hrd, hzd, rdmax, zdmax, 
				      lambda.size(), n_tol, stdev_tol, maxlevel);
  boost::shared_ptr<adaptive_grid<T_grid::T_data> > gg(new adaptive_grid<T_grid::T_data>(-grid_extent, grid_extent, grid_n, factory));;
  T_grid gr (gg, units);
  
  // build stellar grid
  cout << "\nBuilding stellar emission grid" << endl;
  blackbody bb(3000,1);
  T_lambda rho0s(lambda.size());
  rho0s=bb.emission(lambda);

  typedef emission<polychromatic_policy, T_rng_policy> T_emission;
  auto_ptr<T_emission> em;
  T_lambda long_stdev_tol(lambda.size()); long_stdev_tol = stdev_tol(0);
  T_lambda long_n_tol(lambda.size()); long_n_tol=1e300;
  validation_factory<stellar_grid> emfactory (rho0s, hrs, hzs, rsmax, zsmax, 
					      lambda.size(), long_n_tol,
					      long_stdev_tol, maxlevel);
  stellar_grid emgr (-grid_extent, grid_extent, grid_n, emfactory);
  vector<emission<polychromatic_policy, T_rng_policy>* > emitters;
  array_1 L_bol(emgr.n_cells()); L_bol=0;
  int cc=0;
  for (stellar_grid::iterator c=emgr.begin(); c!=emgr.end(); ++c, ++cc) {
    // CRUCIAL: restore the cell pointers for the grid_cell_emission objects
    // this is a flaw in the design of the factory and grid_cell_emission object
    const stellar_grid::T_cell& thec=*c;
    c->data()->set_cell(c);
    emitters.push_back(c->data());
    //c->data()->get_emission()=emfactory.get_rho(*c)*c->volume();
    c->data()->get_emission() *= c->volume();
    const T_float L=integrate_quantity(c->data()->get_emission(), lambda, false);
    L_bol(cc)=L;
    c->data()->get_emission()/=L;
  }
  cout << "Total L_bol " << sum(L_bol) << endl;

  em.reset(new emission_collection<polychromatic_policy, cumulative_sampling, T_rng_policy> (emitters, L_bol));
  
  //cout << bb.emission(lambda); 
  cout << "\nSetting up emergence" << endl;
  std::vector<std::pair<T_float, T_float> > cam_pos;
  const T_float phi=0.3;
  cam_pos.push_back(std:: make_pair (0.0, phi) );
  cam_pos.push_back(std:: make_pair (60./180*pi, phi) );
  cam_pos.push_back(std:: make_pair (80./180*pi, phi) );
  cam_pos.push_back(std:: make_pair (87./180*pi, phi) );
  cam_pos.push_back(std:: make_pair (90./180*pi, phi) );
  auto_ptr<T_emergence> e(new T_emergence(cam_pos, cam_dist, cam_fov, npix));

  cout << "shooting" << endl;

  std::vector<T_rng::T_state> states= generate_random_states (n_threads);
  T_rng_policy rng;

  xfer<T_dust_model, T_grid, T_rng_policy> x (gr, *em, *e, model, rng, 1e-2, 0,n_forced,10., false, true);
  long n_rays_actual = 0;
  dummy_terminator t;
  T_biaser b(vis_ref_lambda);
  shoot (x, scatter_shooter<T_biaser> (b), t, states, n_rays,
  	 n_rays_actual, n_threads,0, "");
  T_float normalization = L_tot/n_rays_actual;

  FITS output(output_file, Write);
  e->write_images(output,normalization,units, "STAR", false, false);

  // now we start the interesting stuff...
  // make dust mass vector
  array_2 dust_mass(gr.n_cells(), model.n_scatterers());
  array_2 star_intensity(gr.intensities().shape());

  cc=0;
  const T_float area_factor = gr.area_factor();
  for(T_grid::const_iterator c=gr.begin(); c!=gr.end(); ++c, ++cc) {
    dust_mass(cc, Range::all()) = c->data()->densities()*c->volume();
    star_intensity(cc, Range::all()) =
      (gr.intensities()(cc,Range::all())*normalization/
       (4*constants::pi*c->volume()*area_factor));
  }

  boost::shared_ptr<adaptive_grid<T_ir_grid::T_data> > ggg(new adaptive_grid<T_ir_grid::T_data>(-grid_extent, grid_extent, grid_n, gg->get_structure()));
  
  T_ir_grid irg(ggg, dust_mass, lambda, units);

  array_2 dust_intensity(gr.n_cells(), lambda.size());
  dust_intensity=0;

  T_biaser b2(ir_ref_lambda);
  scatter_shooter<T_biaser> s(b2);

  // Determine dust equilibrium
  determine_dust_equilibrium_intensities
    (model,
     gr,
     irg,
     lambda,
     star_intensity,
     dust_intensity,
     states,
     n_rays_temp,
     n_threads,
     rel_tol,
     abs_tol,
     1e-2,
     1,
     10.,
     s,
     t,
     false);

  // now generate output
  std::vector<grain_model<polychromatic_scatterer_policy, 
    mcrx_rng_policy>*> grain_models;
  for(int i=0; i< model.n_scatterers(); ++i) {
    grain_models.push_back
      (dynamic_cast<grain_model<polychromatic_scatterer_policy, 
       mcrx_rng_policy>*>(&model.get_scatterer(i)));
    grain_models.back()->resample(lambda);
  }

  e.reset(new T_emergence(cam_pos, cam_dist, cam_fov, npix));
  irg.calculate_SED(array_2(star_intensity + dust_intensity), grain_models,
		    t);
  gr.reset();
  xfer<T_dust_model, T_grid, T_rng_policy> x2 (gr, irg, *e, model, rng, 1e-2, 0,0,10.);

  n_rays_actual = 0;
  // with identically zero dust luminosity we get a crash...
  if(irg.total_emission_weight()>0) {
    shoot (x2, s, t, states, n_rays, n_rays_actual, n_threads,0, "");
    normalization = irg.total_emission_weight()/n_rays_actual;
    e->write_images(output,normalization,units, "IR", false, false);
  }
}

int main (int, char**)
{
  baes_validation(0,"baes_tau0.fits");
  baes_validation(0.5,"baes_tau0.5.fits");
  baes_validation(1,"baes_tau1.fits");
  baes_validation(5,"baes_tau5.fits");
}
